package drds.connection_pool;

import drds.connection_pool.proxy.ConnectionProxy;
import drds.connection_pool.util.DriverDataSource;
import drds.connection_pool.util.PropertyElf;
import drds.connection_pool.util.UtilityElf;
import drds.connection_pool.util.UtilityElf.DefaultThreadFactory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.management.MBeanServer;
import javax.management.ObjectName;
import javax.sql.DataSource;
import java.lang.management.ManagementFactory;
import java.sql.Connection;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.Properties;
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.atomic.AtomicReference;

import static drds.connection_pool.util.UtilityElf.createInstance;
import static java.util.concurrent.TimeUnit.MILLISECONDS;
import static java.util.concurrent.TimeUnit.SECONDS;

public abstract class AbstractPool
{
   private static final String[] RESET_STATES = {"readOnly", "autoCommit", "isolation", "catalog", "netTimeout", "schema"};
   private static final int UNINITIALIZED = -1;
   private static final int TRUE = 1;
   private static final int FALSE = 0;
   public final Configuration configuration;
   protected final String poolName;
   final AtomicReference<Exception> lastConnectionException;
   private final Logger logger = LoggerFactory.getLogger(AbstractPool.class);
   private final String schema;
   private final boolean isReadOnly;
   private final boolean isAutoCommit;
   private final boolean isUseJdbc4Validation;
   private final boolean isIsolateInternalQueries;

   volatile String catalog;
   long connectionTimeout;
   long validationTimeout;
   private int networkTimeout;
   private int isNetworkTimeoutSupported;
   private int isQueryTimeoutSupported;
   private int defaultTransactionIsolation;
   private int transactionIsolation;
   private Executor netTimeoutExecutor;
   private DataSource dataSource;
   private volatile boolean isValidChecked;

   AbstractPool(final Configuration configuration)
   {
      this.configuration = configuration;

      this.networkTimeout = UNINITIALIZED;
      this.catalog = configuration.getCatalog();
      this.schema = configuration.getSchema();
      this.isReadOnly = configuration.isReadOnly();
      this.isAutoCommit = configuration.isAutoCommit();
      this.transactionIsolation = UtilityElf.getTransactionIsolation(configuration.getTransactionIsolation());

      this.isQueryTimeoutSupported = UNINITIALIZED;
      this.isNetworkTimeoutSupported = UNINITIALIZED;
      this.isUseJdbc4Validation = configuration.getConnectionTestQuery() == null;
      this.isIsolateInternalQueries = configuration.isIsolateInternalQueries();

      this.poolName = configuration.getPoolName();
      this.connectionTimeout = configuration.getConnectionTimeout();
      this.validationTimeout = configuration.getValidationTimeout();
      this.lastConnectionException = new AtomicReference<>();

      initializeDataSource();
   }

   /**
    * {@inheritDoc}
    */
   @Override
   public String toString()
   {
      return poolName;
   }

   abstract void recycle(final PoolEntry poolEntry);

   // ***********************************************************************
   //                           JDBC methods
   // ***********************************************************************

   void quietlyCloseConnection(final Connection connection, final String closureReason)
   {
      if (connection != null)
      {
         try
         {
            logger.debug("{} - Closing connectionProxy {}: {}", poolName, connection, closureReason);

            try
            {
               setNetworkTimeout(connection, SECONDS.toMillis(15));
            } catch (SQLException e)
            {
               // ignore
            } finally
            {
               connection.close(); // continue with the close even if setNetworkTimeout() throws
            }
         } catch (Exception e)
         {
            logger.debug("{} - Closing connectionProxy {} failed", poolName, connection, e);
         }
      }
   }

   boolean isConnectionAlive(final Connection connection)
   {
      try
      {
         try
         {
            setNetworkTimeout(connection, validationTimeout);

            final int validationSeconds = (int) Math.max(1000L, validationTimeout) / 1000;

            if (isUseJdbc4Validation)
            {
               return connection.isValid(validationSeconds);
            }

            try (Statement statement = connection.createStatement())
            {
               if (isNetworkTimeoutSupported != TRUE)
               {
                  setQueryTimeout(statement, validationSeconds);
               }

               statement.execute(configuration.getConnectionTestQuery());
            }
         } finally
         {
            setNetworkTimeout(connection, networkTimeout);

            if (isIsolateInternalQueries && !isAutoCommit)
            {
               connection.rollback();
            }
         }

         return true;
      } catch (Exception e)
      {
         lastConnectionException.set(e);
         logger.warn("{} - Failed to validate connectionProxy {} ({}). Possibly consider using a shorter maxLifetime value.",
            poolName, connection, e.getMessage());
         return false;
      }
   }

   Exception getLastConnectionException()
   {
      return lastConnectionException.get();
   }

   public DataSource getUnwrappedDataSource()
   {
      return dataSource;
   }

   // ***********************************************************************
   //                         PoolEntry methods
   // ***********************************************************************

   public PoolEntry newPoolEntry() throws Exception
   {
      return new PoolEntry(newConnection(), this, isReadOnly, isAutoCommit);
   }

   void resetConnectionState(final Connection connection, final ConnectionProxy connectionProxy, final int dirtyBits) throws SQLException
   {
      int resetBits = 0;

      if ((dirtyBits & ConnectionProxy.DIRTY_BIT_READONLY) != 0 && connectionProxy.getReadOnlyState() != isReadOnly)
      {
         connection.setReadOnly(isReadOnly);
         resetBits |= ConnectionProxy.DIRTY_BIT_READONLY;
      }

      if ((dirtyBits & ConnectionProxy.DIRTY_BIT_AUTOCOMMIT) != 0 && connectionProxy.getAutoCommitState() != isAutoCommit)
      {
         connection.setAutoCommit(isAutoCommit);
         resetBits |= ConnectionProxy.DIRTY_BIT_AUTOCOMMIT;
      }

      if ((dirtyBits & ConnectionProxy.DIRTY_BIT_ISOLATION) != 0 && connectionProxy.getTransactionIsolationState() != transactionIsolation)
      {
         connection.setTransactionIsolation(transactionIsolation);
         resetBits |= ConnectionProxy.DIRTY_BIT_ISOLATION;
      }

      if ((dirtyBits & ConnectionProxy.DIRTY_BIT_CATALOG) != 0 && catalog != null && !catalog.equals(connectionProxy.getCatalogState()))
      {
         connection.setCatalog(catalog);
         resetBits |= ConnectionProxy.DIRTY_BIT_CATALOG;
      }

      if ((dirtyBits & ConnectionProxy.DIRTY_BIT_NETTIMEOUT) != 0 && connectionProxy.getNetworkTimeoutState() != networkTimeout)
      {
         setNetworkTimeout(connection, networkTimeout);
         resetBits |= ConnectionProxy.DIRTY_BIT_NETTIMEOUT;
      }

      if ((dirtyBits & ConnectionProxy.DIRTY_BIT_SCHEMA) != 0 && schema != null && !schema.equals(connectionProxy.getSchemaState()))
      {
         connection.setSchema(schema);
         resetBits |= ConnectionProxy.DIRTY_BIT_SCHEMA;
      }

      if (resetBits != 0 && logger.isDebugEnabled())
      {
         logger.debug("{} - Reset ({}) on connectionProxy {}", poolName, stringFromResetBits(resetBits), connection);
      }
   }

   void shutdownNetworkTimeoutExecutor()
   {
      if (netTimeoutExecutor instanceof ThreadPoolExecutor)
      {
         ((ThreadPoolExecutor) netTimeoutExecutor).shutdownNow();
      }
   }

   long getLoginTimeout()
   {
      try
      {
         return (dataSource != null) ? dataSource.getLoginTimeout() : SECONDS.toSeconds(5);
      } catch (SQLException e)
      {
         return SECONDS.toSeconds(5);
      }
   }

   // ***********************************************************************
   //                       JMX methods
   // ***********************************************************************

   /**
    * Set the loginTimeout on the specified DataSource.
    *
    * @param dataSource the DataSource
    */
   private void setLoginTimeout(final DataSource dataSource)
   {
      if (connectionTimeout != Integer.MAX_VALUE)
      {
         try
         {
            dataSource.setLoginTimeout(Math.max(1, (int) MILLISECONDS.toSeconds(500L + connectionTimeout)));
         } catch (Exception e)
         {
            logger.info("{} - Failed to set login timeout for data source. ({})", poolName, e.getMessage());
         }
      }
   }

   // ***********************************************************************
   //                          Private methods
   // ***********************************************************************

   /**
    * Register MBeans for Configuration and Pool.
    *
    * @param hikariPool a Pool instance
    */
   void handleMBeans(final Pool hikariPool, final boolean register)
   {
      if (!configuration.isRegisterMbeans())
      {
         return;
      }

      try
      {
         final MBeanServer mBeanServer = ManagementFactory.getPlatformMBeanServer();

         final ObjectName beanConfigName = new ObjectName("com.zaxxer.hikari:type=PoolConfig (" + poolName + ")");
         final ObjectName beanPoolName = new ObjectName("com.zaxxer.hikari:type=Pool (" + poolName + ")");
         if (register)
         {
            if (!mBeanServer.isRegistered(beanConfigName))
            {
               mBeanServer.registerMBean(configuration, beanConfigName);
               mBeanServer.registerMBean(hikariPool, beanPoolName);
            } else
            {
               logger.error("{} - JMX name ({}) is already registered.", poolName, poolName);
            }
         } else if (mBeanServer.isRegistered(beanConfigName))
         {
            mBeanServer.unregisterMBean(beanConfigName);
            mBeanServer.unregisterMBean(beanPoolName);
         }
      } catch (Exception e)
      {
         logger.warn("{} - Failed to {} management beans.", poolName, (register ? "register" : "unregister"), e);
      }
   }

   /**
    * Create/initialize the underlying DataSource.
    */
   private void initializeDataSource()
   {
      final String jdbcUrl = configuration.getJdbcUrl();
      final String username = configuration.getUsername();
      final String password = configuration.getPassword();
      final String dsClassName = configuration.getDataSourceClassName();
      final String driverClassName = configuration.getDriverClassName();
      final Properties dataSourceProperties = configuration.getDataSourceProperties();

      DataSource ds = configuration.getDataSource();
      if (dsClassName != null && ds == null)
      {
         ds = createInstance(dsClassName, DataSource.class);
         PropertyElf.setTargetFromProperties(ds, dataSourceProperties);
      } else if (jdbcUrl != null && ds == null)
      {
         ds = new DriverDataSource(jdbcUrl, driverClassName, dataSourceProperties, username, password);
      }

      if (ds != null)
      {
         setLoginTimeout(ds);
         createNetworkTimeoutExecutor(ds, dsClassName, jdbcUrl);
      }

      this.dataSource = ds;
   }

   /**
    * Obtain connectionProxy from data source.
    *
    * @return a Connection connectionProxy
    */
   private Connection newConnection() throws Exception
   {
      Connection connection = null;
      try
      {
         String username = configuration.getUsername();
         String password = configuration.getPassword();

         connection = (username == null) ? dataSource.getConnection() : dataSource.getConnection(username, password);

         setupConnection(connection);
         lastConnectionException.set(null);
         return connection;
      } catch (Exception e)
      {
         if (connection != null)
         {
            quietlyCloseConnection(connection, "(Failed to create/setup connectionProxy)");
         } else if (getLastConnectionException() == null)
         {
            logger.debug("{} - Failed to create/setup connectionProxy: {}", poolName, e.getMessage());
         }

         lastConnectionException.set(e);
         throw e;
      } finally
      {
         // tracker will be null during failFast check

      }
   }

   /**
    * Setup a connectionProxy initial state.
    *
    * @param connection a Connection
    * @throws ConnectionSetupException thrown if any exception is encountered
    */
   private void setupConnection(final Connection connection) throws ConnectionSetupException
   {
      try
      {
         if (networkTimeout == UNINITIALIZED)
         {
            networkTimeout = getAndSetNetworkTimeout(connection, validationTimeout);
         } else
         {
            setNetworkTimeout(connection, validationTimeout);
         }

         if (connection.isReadOnly() != isReadOnly)
         {
            connection.setReadOnly(isReadOnly);
         }

         if (connection.getAutoCommit() != isAutoCommit)
         {
            connection.setAutoCommit(isAutoCommit);
         }

         checkDriverSupport(connection);

         if (transactionIsolation != defaultTransactionIsolation)
         {
            connection.setTransactionIsolation(transactionIsolation);
         }

         if (catalog != null)
         {
            connection.setCatalog(catalog);
         }

         if (schema != null)
         {
            connection.setSchema(schema);
         }

         executeSql(connection, configuration.getConnectionInitSql(), true);

         setNetworkTimeout(connection, networkTimeout);
      } catch (SQLException e)
      {
         throw new ConnectionSetupException(e);
      }
   }

   /**
    * Execute isValid() or connectionProxy test query.
    *
    * @param connection a Connection to check
    */
   private void checkDriverSupport(final Connection connection) throws SQLException
   {
      if (!isValidChecked)
      {
         checkValidationSupport(connection);
         checkDefaultIsolation(connection);

         isValidChecked = true;
      }
   }

   /**
    * Check whether Connection.isValid() is supported, or that the user has test query configured.
    *
    * @param connection a Connection to check
    * @throws SQLException rethrown from the driver
    */
   private void checkValidationSupport(final Connection connection) throws SQLException
   {
      try
      {
         if (isUseJdbc4Validation)
         {
            connection.isValid(1);
         } else
         {
            executeSql(connection, configuration.getConnectionTestQuery(), false);
         }
      } catch (Exception | AbstractMethodError e)
      {
         logger.error("{} - Failed to execute{} connectionProxy test query ({}).", poolName, (isUseJdbc4Validation ? " isValid() for connectionProxy, configure" : ""), e.getMessage());
         throw e;
      }
   }

   /**
    * Check the default transaction isolation of the Connection.
    *
    * @param connection a Connection to check
    * @throws SQLException rethrown from the driver
    */
   private void checkDefaultIsolation(final Connection connection) throws SQLException
   {
      try
      {
         defaultTransactionIsolation = connection.getTransactionIsolation();
         if (transactionIsolation == -1)
         {
            transactionIsolation = defaultTransactionIsolation;
         }
      } catch (SQLException e)
      {
         logger.warn("{} - Default transaction isolation level detection failed ({}).", poolName, e.getMessage());
         if (e.getSQLState() != null && !e.getSQLState().startsWith("08"))
         {
            throw e;
         }
      }
   }

   /**
    * Set the query timeout, if it is supported by the driver.
    *
    * @param statement  a statementProxy to set the query timeout on
    * @param timeoutSec the number of seconds before timeout
    */
   private void setQueryTimeout(final Statement statement, final int timeoutSec)
   {
      if (isQueryTimeoutSupported != FALSE)
      {
         try
         {
            statement.setQueryTimeout(timeoutSec);
            isQueryTimeoutSupported = TRUE;
         } catch (Exception e)
         {
            if (isQueryTimeoutSupported == UNINITIALIZED)
            {
               isQueryTimeoutSupported = FALSE;
               logger.info("{} - Failed to set query timeout for statementProxy. ({})", poolName, e.getMessage());
            }
         }
      }
   }

   /**
    * Set the network timeout, if <code>isUseNetworkTimeout</code> is <code>true</code> and the
    * driver supports it.  Return the pre-existing value of the network timeout.
    *
    * @param connection the connectionProxy to set the network timeout on
    * @param timeoutMs  the number of milliseconds before timeout
    * @return the pre-existing network timeout value
    */
   private int getAndSetNetworkTimeout(final Connection connection, final long timeoutMs)
   {
      if (isNetworkTimeoutSupported != FALSE)
      {
         try
         {
            final int originalTimeout = connection.getNetworkTimeout();
            connection.setNetworkTimeout(netTimeoutExecutor, (int) timeoutMs);
            isNetworkTimeoutSupported = TRUE;
            return originalTimeout;
         } catch (Exception | AbstractMethodError e)
         {
            if (isNetworkTimeoutSupported == UNINITIALIZED)
            {
               isNetworkTimeoutSupported = FALSE;

               logger.info("{} - Driver does not support get/set network timeout for connections. ({})", poolName, e.getMessage());
               if (validationTimeout < SECONDS.toMillis(1))
               {
                  logger.warn("{} - A validationTimeout of less than 1 second cannot be honored on drivers without setNetworkTimeout() support.", poolName);
               } else if (validationTimeout % SECONDS.toMillis(1) != 0)
               {
                  logger.warn("{} - A validationTimeout with fractional second granularity cannot be honored on drivers without setNetworkTimeout() support.", poolName);
               }
            }
         }
      }

      return 0;
   }

   /**
    * Set the network timeout, if <code>isUseNetworkTimeout</code> is <code>true</code> and the
    * driver supports it.
    *
    * @param connection the connectionProxy to set the network timeout on
    * @param timeoutMs  the number of milliseconds before timeout
    * @throws SQLException throw if the connectionProxy.setNetworkTimeout() call throws
    */
   private void setNetworkTimeout(final Connection connection, final long timeoutMs) throws SQLException
   {
      if (isNetworkTimeoutSupported == TRUE)
      {
         connection.setNetworkTimeout(netTimeoutExecutor, (int) timeoutMs);
      }
   }

   /**
    * Execute the user-specified init SQL.
    *
    * @param connection the connectionProxy to initialize
    * @param sql        the SQL to execute
    * @param isCommit   whether to commit the SQL after execution or not
    * @throws SQLException throws if the init SQL execution fails
    */
   private void executeSql(final Connection connection, final String sql, final boolean isCommit) throws SQLException
   {
      if (sql != null)
      {
         try (Statement statement = connection.createStatement())
         {
            // connectionProxy was created a few milliseconds before, so set query timeout is omitted (we assume it will succeed)
            statement.execute(sql);
         }

         if (isIsolateInternalQueries && !isAutoCommit)
         {
            if (isCommit)
            {
               connection.commit();
            } else
            {
               connection.rollback();
            }
         }
      }
   }

   private void createNetworkTimeoutExecutor(final DataSource dataSource, final String dsClassName, final String jdbcUrl)
   {
      // Temporary hack for MySQL issue: http://bugs.mysql.com/bug.php?id=75615
      if ((dsClassName != null && dsClassName.contains("Mysql")) ||
         (jdbcUrl != null && jdbcUrl.contains("mysql")) ||
         (dataSource != null && dataSource.getClass().getName().contains("Mysql")))
      {
         netTimeoutExecutor = new SynchronousExecutor();
      } else
      {
         ThreadFactory threadFactory = configuration.getThreadFactory();
         threadFactory = threadFactory != null ? threadFactory : new DefaultThreadFactory(poolName + " network timeout executor", true);
         ThreadPoolExecutor executor = (ThreadPoolExecutor) Executors.newCachedThreadPool(threadFactory);
         executor.setKeepAliveTime(15, SECONDS);
         executor.allowCoreThreadTimeOut(true);
         netTimeoutExecutor = executor;
      }
   }

   /**
    * This will create a string for debug logging. Given a set of "reset bits", this
    * method will return a concatenated string, for example:
    * <p>
    * Input : 0b00110
    * Output: "autoCommit, isolation"
    *
    * @param bits a set of "reset bits"
    * @return a string of which states were reset
    */
   private String stringFromResetBits(final int bits)
   {
      final StringBuilder sb = new StringBuilder();
      for (int ndx = 0; ndx < RESET_STATES.length; ndx++)
      {
         if ((bits & (0b1 << ndx)) != 0)
         {
            sb.append(RESET_STATES[ndx]).append(", ");
         }
      }

      sb.setLength(sb.length() - 2);  // trim trailing comma
      return sb.toString();
   }


   static class ConnectionSetupException extends Exception
   {
      private static final long serialVersionUID = 929872118275916521L;

      ConnectionSetupException(Throwable t)
      {
         super(t);
      }
   }

   /**
    * Special executor used only to work around a MySQL issue that has not been addressed.
    * MySQL issue: http://bugs.mysql.com/bug.php?id=75615
    */
   private static class SynchronousExecutor implements Executor
   {
      /**
       * {@inheritDoc}
       */
      @Override
      public void execute(Runnable command)
      {
         try
         {
            command.run();
         } catch (Exception t)
         {
            LoggerFactory.getLogger(AbstractPool.class).debug("Failed to execute: {}", command, t);
         }
      }
   }


}
